
# import PyPDF2
from pdf2image import convert_from_path
# from PIL import Image

from extra_service.llm import GlobalLLMClient
from extra_service.llm import *

prompt = """图片是一篇论文的第一页，请你从中提取这篇论文所属的机构是什么？
如果有多个机构，请用逗号分隔。对于明显的人名、地名和国家等，可以忽略, 请只提取机构名称。
请你提取:
"""

def Ask(image_path:str):
    c = Conversation()
    c.Add(Content(
        role="user",
        content=[
            ContentImage(image_path),
            ContentText(prompt)
        ]
    ))
    result , c = GlobalLLMClient.Predict(c)
    print(result)
    # print(c)

    # 2
    c.Add(Content(
        role="user",
        content=[
            ContentText(f"上次提取机构的名称的结果是：{result}"),
            ContentText("上面是提取得到的机构，请你将这些机构每行一个写出来:"),
        ]
    ))
    result , c = GlobalLLMClient.Predict(c)
    # print("=================\n")
    # print(c)
    print(result)

    # 3
    all = result.split("\n")

    # 去除所有数字
    def noNumber(s:str) -> str:
        return ''.join([i for i in s if not i.isdigit()])

    all = [noNumber(i) for i in all]
    all = [i.strip() for i in all]
    return all # 返回所有机构名称


def extract_first_half_image(pdf_path, output_image_path):
    # 使用pdf2image将PDF第一页转换为图像
    images = convert_from_path(pdf_path, first_page=1, last_page=1)
    
    if images:
        first_page_image = images[0]

        # 获取图像尺寸
        width, height = first_page_image.size
        
        # 计算前半部分的区域 (上半部分)
        upper_half = first_page_image.crop((0, 0, width, height // 2))
        
        # 保存前半部分的图像
        upper_half.save(output_image_path)
        
        print(f"Saved upper half of the first page as {output_image_path}")
    else:
        print("Could not convert PDF to image.")


def AskLLMforInstitution(pdf_path: str):
    # 1. 将 PDF 的第一页转换为图像
    output_image_path = 'output_image.png'
    extract_first_half_image(pdf_path, output_image_path)
    # 2. 询问 LLM 模型
    all = Ask(output_image_path)
    return all
